import numpy as np

eps = 1  # 1e-38?


def make_fourier_showable(f):
    '''
    Process fourier transformation result to let it be suit for showing. (magnitude, log)
    :param f: fourier transformation result
    :return: the data suitable for showing
    '''
    f_show = np.log(cv.magnitude(f[:, :, 0], f[:, :, 1]) + eps)
    # f_show = cv.normalize(f_show, None, 0., 1., norm_type=cv.NORM_MINMAX)
    return f_show


if '__main__' == __name__:
    from python_ai.common.xcommon import *
    import cv2 as cv
    import matplotlib.pyplot as plt
    import os

    # load img
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    path = '../../../../../../large_data/pic/messi5.jpg'
    IMG_PATH = os.path.join(BASE_DIR, path)
    img = cv.imread(IMG_PATH, cv.IMREAD_GRAYSCALE)
    print_numpy_ndarray_info(img, 'img')

    # size of ori img
    H, W = img.shape
    H2 = (H - 1) // 2
    W2 = (W - 1) // 2

    # plt config
    plt.figure(figsize=[12, 8])
    plt_info = dict(spr=2, spc=3, spn=0)

    # show ori img
    show_cv_img_by_plt(plt_info, img, 'ori', cmap='gray')

    # fourier transformation
    # f: fourier
    img_flt32 = img.astype(np.float32)
    f = cv.dft(img_flt32, flags=cv.DFT_COMPLEX_OUTPUT)
    print_numpy_ndarray_info(f, 'f')
    # f_show: fourier for showing
    f_show = make_fourier_showable(f)
    print_numpy_ndarray_info(f_show, 'f_show')
    show_cv_img_by_plt(plt_info, f_show, 'f_show', cmap='gray')

    # shift fourier (shift 4 corners to center)
    # f_shift: shifted fourier
    f_shift = np.fft.fftshift(f)
    print_numpy_ndarray_info(f_shift, 'f_shift')
    # f_shift_show: shifted fourier for showing
    f_shift_show = make_fourier_showable(f_shift)
    print_numpy_ndarray_info(f_shift_show, 'f_shift_show')
    show_cv_img_by_plt(plt_info, f_shift_show, 'f_shift_show', cmap='gray')

    # high pass: erase center
    HALF_RANGE = 60
    f_shift[H2 - HALF_RANGE:H2 + HALF_RANGE, W2 - HALF_RANGE:W2 + HALF_RANGE, :] = 0
    print_numpy_ndarray_info(f_shift, 'f_shift (high pass)')
    f_shift_show = make_fourier_showable(f_shift)
    print_numpy_ndarray_info(f_shift_show, 'f_shift_show (high pass)')
    show_cv_img_by_plt(plt_info, f_shift_show, 'f_shift_show (high pass)', cmap='gray')

    # inverse shift
    # f: inverse shifted: the fourier
    f = np.fft.ifftshift(f_shift)
    print_numpy_ndarray_info(f, 'f (high pass)')
    # f_show: inverse shifted: the fourier for showing
    f_show = make_fourier_showable(f)
    print_numpy_ndarray_info(f_show, 'f_show (high pass)')
    show_cv_img_by_plt(plt_info, f_show, 'f_show (high pass)', cmap='gray')

    # inverse fourier
    img_complex = cv.idft(f)
    print_numpy_ndarray_info(img_complex, 'img_complex')
    img_only_magnitude = cv.magnitude(img_complex[:, :, 0], img_complex[:, :, 1])
    print_numpy_ndarray_info(img_only_magnitude, 'img_only_magnitude')
    img_mag_log = np.log(img_only_magnitude + eps)
    print_numpy_ndarray_info(img_mag_log, 'img_mag_log')
    # img = make_fourier_showable(img_complex)
    # print_numpy_ndarray_info(img, 'img')
    show_cv_img_by_plt(plt_info, img_only_magnitude, 'img (high pass)', cmap='gray')  # do not log
    img_cv = cv.normalize(img_only_magnitude, None, 0., 1., cv.NORM_MINMAX)
    img_cv = imzoom_rect(img_cv, (800, 600))
    cv.imshow('high pass', img_cv)

    plt.show()
    cv.waitKey(0)
